{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 1: Install Astra Assistants"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-01T04:09:15.962971Z",
     "start_time": "2024-07-01T04:09:14.225767Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "kgFdcp7fuQwO",
    "outputId": "61831af6-377d-492e-a1c4-91ebae2c1145"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: astra-assistants in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (2.0.5)\n",
      "Requirement already satisfied: gradio in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (4.21.0)\n",
      "Requirement already satisfied: aiohttp<4.0.0,>=3.9.4 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from astra-assistants) (3.9.5)\n",
      "Requirement already satisfied: boto3<2.0.0,>=1.34.31 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from astra-assistants) (1.34.86)\n",
      "Requirement already satisfied: httpx<0.27.0,>=0.26.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from astra-assistants) (0.26.0)\n",
      "Requirement already satisfied: litellm<2.0.0,>=1.36.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from astra-assistants) (1.40.29)\n",
      "Requirement already satisfied: openai<2.0.0,>=1.20.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from astra-assistants) (1.35.9)\n",
      "Requirement already satisfied: python-dotenv<2.0.0,>=1.0.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from astra-assistants) (1.0.1)\n",
      "Requirement already satisfied: aiofiles<24.0,>=22.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (23.2.1)\n",
      "Requirement already satisfied: altair<6.0,>=4.2.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (5.2.0)\n",
      "Requirement already satisfied: fastapi in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.110.0)\n",
      "Requirement already satisfied: ffmpy in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.3.2)\n",
      "Requirement already satisfied: gradio-client==0.12.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.12.0)\n",
      "Requirement already satisfied: huggingface-hub>=0.19.3 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.21.4)\n",
      "Requirement already satisfied: importlib-resources<7.0,>=1.3 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (6.3.1)\n",
      "Requirement already satisfied: jinja2<4.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (3.1.3)\n",
      "Requirement already satisfied: markupsafe~=2.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (2.1.3)\n",
      "Requirement already satisfied: matplotlib~=3.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (3.8.3)\n",
      "Requirement already satisfied: numpy~=1.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (1.26.4)\n",
      "Requirement already satisfied: orjson~=3.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (3.9.10)\n",
      "Requirement already satisfied: packaging in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (23.2)\n",
      "Requirement already satisfied: pandas<3.0,>=1.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (2.2.1)\n",
      "Requirement already satisfied: pillow<11.0,>=8.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (10.2.0)\n",
      "Requirement already satisfied: pydantic>=2.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (2.7.0)\n",
      "Requirement already satisfied: pydub in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.25.1)\n",
      "Requirement already satisfied: python-multipart>=0.0.9 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.0.9)\n",
      "Requirement already satisfied: pyyaml<7.0,>=5.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (6.0.1)\n",
      "Requirement already satisfied: ruff>=0.2.2 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.3.3)\n",
      "Requirement already satisfied: semantic-version~=2.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (2.10.0)\n",
      "Requirement already satisfied: tomlkit==0.12.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.12.0)\n",
      "Requirement already satisfied: typer<1.0,>=0.9 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from typer[all]<1.0,>=0.9->gradio) (0.9.0)\n",
      "Requirement already satisfied: typing-extensions~=4.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (4.9.0)\n",
      "Requirement already satisfied: uvicorn>=0.14.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio) (0.28.0)\n",
      "Requirement already satisfied: fsspec in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio-client==0.12.0->gradio) (2024.3.1)\n",
      "Requirement already satisfied: websockets<12.0,>=10.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from gradio-client==0.12.0->gradio) (11.0.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.4->astra-assistants) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.4->astra-assistants) (23.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.4->astra-assistants) (1.4.1)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.4->astra-assistants) (6.0.5)\n",
      "Requirement already satisfied: yarl<2.0,>=1.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.4->astra-assistants) (1.9.4)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from aiohttp<4.0.0,>=3.9.4->astra-assistants) (4.0.3)\n",
      "Requirement already satisfied: jsonschema>=3.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from altair<6.0,>=4.2.0->gradio) (4.19.2)\n",
      "Requirement already satisfied: toolz in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from altair<6.0,>=4.2.0->gradio) (0.12.1)\n",
      "Requirement already satisfied: botocore<1.35.0,>=1.34.86 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from boto3<2.0.0,>=1.34.31->astra-assistants) (1.34.86)\n",
      "Requirement already satisfied: jmespath<2.0.0,>=0.7.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from boto3<2.0.0,>=1.34.31->astra-assistants) (1.0.1)\n",
      "Requirement already satisfied: s3transfer<0.11.0,>=0.10.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from boto3<2.0.0,>=1.34.31->astra-assistants) (0.10.1)\n",
      "Requirement already satisfied: anyio in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from httpx<0.27.0,>=0.26.0->astra-assistants) (3.7.1)\n",
      "Requirement already satisfied: certifi in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from httpx<0.27.0,>=0.26.0->astra-assistants) (2024.2.2)\n",
      "Requirement already satisfied: httpcore==1.* in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from httpx<0.27.0,>=0.26.0->astra-assistants) (1.0.4)\n",
      "Requirement already satisfied: idna in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from httpx<0.27.0,>=0.26.0->astra-assistants) (3.4)\n",
      "Requirement already satisfied: sniffio in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from httpx<0.27.0,>=0.26.0->astra-assistants) (1.3.0)\n",
      "Requirement already satisfied: h11<0.15,>=0.13 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from httpcore==1.*->httpx<0.27.0,>=0.26.0->astra-assistants) (0.14.0)\n",
      "Requirement already satisfied: filelock in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from huggingface-hub>=0.19.3->gradio) (3.13.1)\n",
      "Requirement already satisfied: requests in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from huggingface-hub>=0.19.3->gradio) (2.31.0)\n",
      "Requirement already satisfied: tqdm>=4.42.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from huggingface-hub>=0.19.3->gradio) (4.66.2)\n",
      "Requirement already satisfied: click in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from litellm<2.0.0,>=1.36.0->astra-assistants) (8.1.7)\n",
      "Requirement already satisfied: ijson in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from litellm<2.0.0,>=1.36.0->astra-assistants) (3.3.0)\n",
      "Requirement already satisfied: importlib-metadata>=6.8.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from litellm<2.0.0,>=1.36.0->astra-assistants) (7.1.0)\n",
      "Requirement already satisfied: tiktoken>=0.7.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from litellm<2.0.0,>=1.36.0->astra-assistants) (0.7.0)\n",
      "Requirement already satisfied: tokenizers in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from litellm<2.0.0,>=1.36.0->astra-assistants) (0.19.1)\n",
      "Requirement already satisfied: contourpy>=1.0.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from matplotlib~=3.0->gradio) (1.2.0)\n",
      "Requirement already satisfied: cycler>=0.10 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from matplotlib~=3.0->gradio) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from matplotlib~=3.0->gradio) (4.50.0)\n",
      "Requirement already satisfied: kiwisolver>=1.3.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from matplotlib~=3.0->gradio) (1.4.5)\n",
      "Requirement already satisfied: pyparsing>=2.3.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from matplotlib~=3.0->gradio) (3.1.2)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from matplotlib~=3.0->gradio) (2.8.2)\n",
      "Requirement already satisfied: distro<2,>=1.7.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from openai<2.0.0,>=1.20.0->astra-assistants) (1.9.0)\n",
      "Requirement already satisfied: pytz>=2020.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from pandas<3.0,>=1.0->gradio) (2023.3.post1)\n",
      "Requirement already satisfied: tzdata>=2022.7 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from pandas<3.0,>=1.0->gradio) (2024.1)\n",
      "Requirement already satisfied: annotated-types>=0.4.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from pydantic>=2.0->gradio) (0.6.0)\n",
      "Requirement already satisfied: pydantic-core==2.18.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from pydantic>=2.0->gradio) (2.18.1)\n",
      "Requirement already satisfied: colorama<0.5.0,>=0.4.3 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from typer[all]<1.0,>=0.9->gradio) (0.4.6)\n",
      "Requirement already satisfied: shellingham<2.0.0,>=1.3.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from typer[all]<1.0,>=0.9->gradio) (1.5.4)\n",
      "Requirement already satisfied: rich<14.0.0,>=10.11.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from typer[all]<1.0,>=0.9->gradio) (13.7.1)\n",
      "Requirement already satisfied: starlette<0.37.0,>=0.36.3 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from fastapi->gradio) (0.36.3)\n",
      "Requirement already satisfied: exceptiongroup in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from anyio->httpx<0.27.0,>=0.26.0->astra-assistants) (1.2.0)\n",
      "Requirement already satisfied: urllib3!=2.2.0,<3,>=1.25.4 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from botocore<1.35.0,>=1.34.86->boto3<2.0.0,>=1.34.31->astra-assistants) (2.1.0)\n",
      "Requirement already satisfied: zipp>=0.5 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from importlib-metadata>=6.8.0->litellm<2.0.0,>=1.36.0->astra-assistants) (3.18.1)\n",
      "Requirement already satisfied: jsonschema-specifications>=2023.03.6 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from jsonschema>=3.0->altair<6.0,>=4.2.0->gradio) (2023.12.1)\n",
      "Requirement already satisfied: referencing>=0.28.4 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from jsonschema>=3.0->altair<6.0,>=4.2.0->gradio) (0.33.0)\n",
      "Requirement already satisfied: rpds-py>=0.7.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from jsonschema>=3.0->altair<6.0,>=4.2.0->gradio) (0.10.6)\n",
      "Requirement already satisfied: six>=1.5 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from python-dateutil>=2.7->matplotlib~=3.0->gradio) (1.16.0)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from requests->huggingface-hub>=0.19.3->gradio) (2.0.4)\n",
      "Requirement already satisfied: markdown-it-py>=2.2.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from rich<14.0.0,>=10.11.0->typer[all]<1.0,>=0.9->gradio) (3.0.0)\n",
      "Requirement already satisfied: pygments<3.0.0,>=2.13.0 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from rich<14.0.0,>=10.11.0->typer[all]<1.0,>=0.9->gradio) (2.15.1)\n",
      "Requirement already satisfied: regex>=2022.1.18 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from tiktoken>=0.7.0->litellm<2.0.0,>=1.36.0->astra-assistants) (2023.12.25)\n",
      "Requirement already satisfied: mdurl~=0.1 in /Users/vrsen/anaconda3/envs/agency-swarm/lib/python3.10/site-packages (from markdown-it-py>=2.2.0->rich<14.0.0,>=10.11.0->typer[all]<1.0,>=0.9->gradio) (0.1.2)\n"
     ]
    }
   ],
   "source": [
    "!pip install astra-assistants gradio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add agency swarm from local\n",
    "import sys\n",
    "sys.path.append('../agency-swarm')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 2: Patch OpenAI Client \n",
    "\n",
    "Before running this step, make sure your Astra DB token is in the .env file.\n",
    "\n",
    "```\n",
    "ASTRA_DB_APPLICATION_TOKEN=AstraCS:...\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Patching OpenAI client, it will now communicate to Astra Assistants API: https://open-assistant-ai.astra.datastax.com/v1/\n",
      "Learn more about Astra at: https://docs.datastax.com/en/astra-db-serverless/tutorials/astra-assistants-api.html\n"
     ]
    }
   ],
   "source": [
    "from openai import OpenAI\n",
    "from astra_assistants import patch\n",
    "from agency_swarm import set_openai_client\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "load_dotenv()\n",
    "\n",
    "client = patch(OpenAI())\n",
    "\n",
    "set_openai_client(client)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 3: Setup Your Agents and Tools  \n",
    "\n",
    "Before running this step, add your model API keys into the .env file using the following format:\n",
    "```\n",
    "PERPLEXITYAI_API_KEY=your_perplexityai_api_key\n",
    "ANTHROPIC_API_KEY=your_anthropic_api_key\n",
    "TOGETHER_API_KEY=your_together_api_key\n",
    "GROQ_API_KEY=your_groq_api_key\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-07-01T04:15:38.532698Z",
     "start_time": "2024-07-01T04:15:26.063677Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 408
    },
    "id": "QxXyyuH7vovu",
    "outputId": "f81a523e-ecad-42ca-955b-c7e53b39d3cd"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Patching OpenAI client, it will now communicate to Astra Assistants API: https://open-assistant-ai.astra.datastax.com/v1/\n",
      "Learn more about Astra at: https://docs.datastax.com/en/astra-db-serverless/tutorials/astra-assistants-api.html\n",
      "Updating agent... CEO\n",
      "Updating agent... Agent2\n"
     ]
    }
   ],
   "source": [
    "from agency_swarm import Agent, Agency\n",
    "from agency_swarm.tools import BaseTool\n",
    "\n",
    "\n",
    "class PrintTool(BaseTool):\n",
    "    \"\"\"\n",
    "    A simple tool that prints input.\n",
    "    \"\"\"\n",
    "    input: str\n",
    "\n",
    "    def run(self):\n",
    "        \"\"\"\n",
    "        This method prints the word 'test'.\n",
    "        \"\"\"\n",
    "        print(self.input)\n",
    "        return f\"{self.input} has been printed.\"\n",
    "\n",
    "ceo = Agent(name=\"CEO\",\n",
    "            description=\"Responsible for client communication, task planning, and management.\",\n",
    "            instructions=\"You must say 'I am using test tool' and then use test tool in the same message.\",\n",
    "            # model=\"perplexity/llama-3-8b-instruct\",\n",
    "            # model=\"anthropic/claude-3-haiku-20240307\",\n",
    "            # model=\"groq/mixtral-8x7b-32768\",\n",
    "            model=\"claude-3-5-sonnet-20240620\",\n",
    "            # model=\"gpt-4o\",\n",
    "            # files_folder=\"./files\",\n",
    "            temperature=0,\n",
    "            tools=[PrintTool])\n",
    "\n",
    "agent2 = Agent(name=\"Agent2\",\n",
    "              description=\"Test agent for demo purposes\",\n",
    "              instructions=\"You are a test agent for demo purposes\",\n",
    "              # files_folder=\"./files\",\n",
    "              model=\"claude-3-5-sonnet-20240620\")\n",
    "\n",
    "agency = Agency([ceo, [ceo, agent2]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Step 4: Run Demo\n",
    "\n",
    "To run gradio, use the special non-streaming Gradio method below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import time\n",
    "\n",
    "\n",
    "def demo_gradio(agency, height=450, dark_mode=True):\n",
    "    \"\"\"\n",
    "    Launches a Gradio-based demo interface for the agency chatbot.\n",
    "\n",
    "    Parameters:\n",
    "        height (int, optional): The height of the chatbot widget in the Gradio interface. Default is 600.\n",
    "        dark_mode (bool, optional): Flag to determine if the interface should be displayed in dark mode. Default is True.\n",
    "        share (bool, optional): Flag to determine if the interface should be shared publicly. Default is False.\n",
    "    This method sets up and runs a Gradio interface, allowing users to interact with the agency's chatbot. It includes a text input for the user's messages and a chatbot interface for displaying the conversation. The method handles user input and chatbot responses, updating the interface dynamically.\n",
    "    \"\"\"\n",
    "    try:\n",
    "        import gradio as gr\n",
    "    except ImportError:\n",
    "        raise Exception(\"Please install gradio: pip install gradio\")\n",
    "\n",
    "    js = \"\"\"function () {\n",
    "      gradioURL = window.location.href\n",
    "      if (!gradioURL.endsWith('?__theme={theme}')) {\n",
    "        window.location.replace(gradioURL + '?__theme={theme}');\n",
    "      }\n",
    "    }\"\"\"\n",
    "\n",
    "    if dark_mode:\n",
    "        js = js.replace(\"{theme}\", \"dark\")\n",
    "    else:\n",
    "        js = js.replace(\"{theme}\", \"light\")\n",
    "\n",
    "    message_file_ids = []\n",
    "    message_file_names = None\n",
    "    recipient_agents = [agent.name for agent in agency.main_recipients]\n",
    "    recipient_agent = agency.main_recipients[0]\n",
    "\n",
    "    with (gr.Blocks(js=js) as demo):\n",
    "        chatbot = gr.Chatbot(height=height)\n",
    "        with gr.Row():\n",
    "            with gr.Column(scale=9):\n",
    "                dropdown = gr.Dropdown(label=\"Recipient Agent\", choices=recipient_agents,\n",
    "                                       value=recipient_agent.name)\n",
    "                msg = gr.Textbox(label=\"Your Message\", lines=4)\n",
    "            with gr.Column(scale=1):\n",
    "                file_upload = gr.Files(label=\"Files\", type=\"filepath\")\n",
    "        button = gr.Button(value=\"Send\", variant=\"primary\")\n",
    "\n",
    "        def handle_dropdown_change(selected_option):\n",
    "            nonlocal recipient_agent\n",
    "            recipient_agent = agency._get_agent_by_name(selected_option)\n",
    "\n",
    "        def handle_file_upload(file_list):\n",
    "            nonlocal message_file_ids\n",
    "            nonlocal message_file_names\n",
    "            message_file_ids = []\n",
    "            message_file_names = []\n",
    "            if file_list:\n",
    "                try:\n",
    "                    for file_obj in file_list:\n",
    "                        with open(file_obj.name, 'rb') as f:\n",
    "                            # Upload the file to OpenAI\n",
    "                            file = agency.main_thread.client.files.create(\n",
    "                                file=f,\n",
    "                                purpose=\"assistants\"\n",
    "                            )\n",
    "                        message_file_ids.append(file.id)\n",
    "                        message_file_names.append(file.filename)\n",
    "                        print(f\"Uploaded file ID: {file.id}\")\n",
    "                    return message_file_ids\n",
    "                except Exception as e:\n",
    "                    print(f\"Error: {e}\")\n",
    "                    return str(e)\n",
    "\n",
    "            return \"No files uploaded\"\n",
    "\n",
    "        def user(user_message, history):\n",
    "            if not user_message:\n",
    "                return user_message, history\n",
    "\n",
    "            if history is None:\n",
    "                history = []\n",
    "\n",
    "            original_user_message = user_message\n",
    "\n",
    "            # Append the user message with a placeholder for bot response\n",
    "            if recipient_agent:\n",
    "                user_message = f\"👤 User @{recipient_agent.name}:\\n\" + user_message.strip()\n",
    "            else:\n",
    "                user_message = f\"👤 User:\" + user_message.strip()\n",
    "\n",
    "            nonlocal message_file_names\n",
    "            if message_file_names:\n",
    "                user_message += \"\\n\\n:paperclip: Files:\\n\" + \"\\n\".join(message_file_names)\n",
    "\n",
    "            return original_user_message, history + [[user_message, None]]\n",
    "\n",
    "        def bot(original_message, history):\n",
    "            nonlocal message_file_ids\n",
    "            nonlocal message_file_names\n",
    "            nonlocal recipient_agent\n",
    "            print(\"Message files: \", message_file_ids)\n",
    "            # Replace this with your actual chatbot logic\n",
    "            gen = agency.get_completion(message=original_message, message_files=message_file_ids,\n",
    "                                        recipient_agent=recipient_agent, yield_messages=True)\n",
    "\n",
    "            message_file_ids = []\n",
    "            message_file_names = []\n",
    "            try:\n",
    "                # Yield each message from the generator\n",
    "                for bot_message in gen:\n",
    "                    if bot_message.sender_name.lower() == \"user\":\n",
    "                        continue\n",
    "\n",
    "                    # sometimes thread stops before bot message is received\n",
    "                    if not bot_message.content:\n",
    "                        main_thread = agency.main_thread\n",
    "                        content = bot_message.content\n",
    "                        num_attempts = 0\n",
    "                        while not content or num_attempts < 30:\n",
    "                            time.sleep(1)\n",
    "                            content = main_thread._get_last_message_text()\n",
    "                            num_attempts += 1\n",
    "\n",
    "                        bot_message.content = content\n",
    "\n",
    "                    message = bot_message.get_formatted_content()\n",
    "\n",
    "                    history.append((None, message))\n",
    "                    yield \"\", history\n",
    "            except StopIteration:\n",
    "                # Handle the end of the conversation if necessary\n",
    "                pass\n",
    "\n",
    "        button.click(\n",
    "            user,\n",
    "            inputs=[msg, chatbot],\n",
    "            outputs=[msg, chatbot]\n",
    "        ).then(\n",
    "            bot, [msg, chatbot], [msg, chatbot]\n",
    "        )\n",
    "        dropdown.change(handle_dropdown_change, dropdown)\n",
    "        file_upload.change(handle_file_upload, file_upload)\n",
    "        msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(\n",
    "            bot, [msg, chatbot], [msg, chatbot]\n",
    "        )\n",
    "\n",
    "        # Enable queuing for streaming intermediate outputs\n",
    "        demo.queue()\n",
    "\n",
    "    # Launch the demo\n",
    "    demo.launch(share=False, debug=True)\n",
    "    return demo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "ExecuteTime": {
     "start_time": "2024-07-01T04:15:38.533935Z"
    },
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "PPedd_AcRXhE",
    "jupyter": {
     "is_executing": true
    },
    "outputId": "7bcb3859-4ae6-4de9-d582-1cdff2ddf9f3"
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Running on local URL:  http://127.0.0.1:7860\n",
      "IMPORTANT: You are using gradio version 4.21.0, however version 4.29.0 is available, please upgrade.\n",
      "--------\n",
      "\n",
      "To create a public link, set `share=True` in `launch()`.\n"
     ]
    },
    {
     "data": {
      "text/html": [
       "<div><iframe src=\"http://127.0.0.1:7860/\" width=\"100%\" height=\"500\" allow=\"autoplay; camera; microphone; clipboard-read; clipboard-write;\" frameborder=\"0\" allowfullscreen></iframe></div>"
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "demo_gradio(agency, height=900)"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
